package org.pombo.gtb;

import java.util.List;

import org.pombo.gtm.Action;
import org.pombo.gtm.ChanceNode;
import org.pombo.gtm.Node;

import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

public class ChanceNodeBuilder extends Builder<ChanceNode> {

	@Override
	public void clear() {
		this.entity = mock(ChanceNode.class);
	}

	public ChanceNodeBuilder withNodeForAction(Node node, Action action) {
		when(this.entity.getNodeForAction(action)).thenReturn(node);
		return this;
	}

	public ChanceNodeBuilder withProbabilityForAction(double probability, Action action) {
		when(this.entity.getProbabilityForAction(action)).thenReturn(probability);
		return this;
	}

	public ChanceNodeBuilder withActions(List<Action> actions) {
		when(this.entity.getActions()).thenReturn(actions);
		return this;
	}

}
